
#include <cmath>
#include "range_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;                                                       
const uint32_t UB_BLOCK_NUM = 100;
const uint32_t MAX_AVAILABLE_UB_BLOCK_NUM = UB_BLOCK_NUM / BUFFER_NUM * BUFFER_NUM;

void TilingParamsCalc(uint32_t length, uint32_t alignNum,
                      uint32_t &tileNum, uint32_t &tileLength, uint32_t &lastTileLength)
{
    uint32_t maxPerCoreElem = alignNum * MAX_AVAILABLE_UB_BLOCK_NUM;  
    tileNum = length / maxPerCoreElem;  
    if ((length % maxPerCoreElem == 0U) || (tileNum == 0U)) 
    {
        if (tileNum == 0U) tileNum = (length + alignNum - 1) / alignNum;                                                                           
        if (length < maxPerCoreElem) 
        {
            uint32_t blockCount = (length + alignNum - 1) / alignNum;  
            blockCount = (blockCount + BUFFER_NUM - 1) / BUFFER_NUM * BUFFER_NUM;                             
            tileLength = blockCount * alignNum; 
            lastTileLength = (tileNum == 1) ? length : (length % alignNum == 0 ? alignNum : length % alignNum);
        } 
        else 
        {
            tileLength = maxPerCoreElem;
            lastTileLength = length - (tileNum - 1) * tileLength;
        }
    } 
    else 
    {
        tileNum++;
        tileLength = maxPerCoreElem;
        lastTileLength = length - (tileNum - 1) * tileLength;
    }
}

static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    printf("[TilingFunc] ENTER\n");
    
    // 检查context指针是否为空
    if (context == nullptr) {
        printf("[TilingFunc] ERROR: TilingContext is null!\n");
        return ge::GRAPH_FAILED;
    }

    uint64_t ubSize;
    uint32_t dataTypeSize = 4; //输出float
    uint32_t alignNum = BLOCK_SIZE / dataTypeSize;

    RangeCustomTilingData tiling;
    
    auto platformInfo = context->GetPlatformInfo();
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    auto coreNum = ascendcPlatform.GetCoreNum();

    auto tensorStart = context->GetInputTensor(0);
    auto tensorEnd = context->GetInputTensor(1);
    auto tensorStep = context->GetInputTensor(2);
    
    if (tensorStart == nullptr || tensorEnd == nullptr || tensorStep == nullptr) {
        printf("[TilingFunc] ERROR: One or more input tensors are null!\n");
        return ge::GRAPH_FAILED;
    }
    
    // 计算Range的输出长度: ceil((end - start) / step)
    float startVal = 0.0f, endVal = 0.0f, stepVal = 1.0f;
    
    // 根据数据类型获取值
    const auto inputDataType_1 = context->GetInputDesc(0)->GetDataType();
    const auto inputDataType_2 = context->GetInputDesc(1)->GetDataType();
    const auto inputDataType_3 = context->GetInputDesc(2)->GetDataType();
    printf("[TilintFunc] initsuccess\n");


    // 获取start值
    switch(inputDataType_1) {
        case ge::DT_FLOAT:
            startVal = *tensorStart->GetData<float>();
            break;
        case ge::DT_INT64:
            startVal = static_cast<float>(*tensorStart->GetData<int64_t>());
            break;
        default:
            printf("[TilingFunc] ERROR: Unsupported start data type\n");
            return ge::GRAPH_FAILED;
    }
    
    // 获取end值
    switch(inputDataType_2) {
        case ge::DT_FLOAT:
            endVal = *tensorEnd->GetData<float>();
            break;
        case ge::DT_INT64:
            endVal = static_cast<float>(*tensorEnd->GetData<int64_t>());
            break;
        default:
            printf("[TilingFunc] ERROR: Unsupported end data type\n");
            return ge::GRAPH_FAILED;
    }
    
    // 获取step值
    switch(inputDataType_3) {
        case ge::DT_FLOAT:
            stepVal = *tensorStep->GetData<float>();
            break;
        case ge::DT_INT64:
            stepVal = static_cast<float>(*tensorStep->GetData<int64_t>());
            break;
        default:
            printf("[TilingFunc] ERROR: Unsupported step data type\n");
            return ge::GRAPH_FAILED;
    }
    printf("[TilingFunc] ENTER COMPUTE:startVal=%.6f, endVal=%.6f, stepVal=%.6f\n", startVal, endVal, stepVal);
    // 计算输出长度: ceil(abs(end - start) / abs(step))
    float diff = std::abs(endVal - startVal);
    float stepAbs = std::abs(stepVal);
    int64_t totalLength = static_cast<int64_t>(std::ceil(diff / stepAbs));
    uint32_t totalLengthAligned = (totalLength % alignNum == 0U)
        ? totalLength
        : ((totalLength + alignNum - 1) / alignNum) * alignNum;

    if ((totalLengthAligned / alignNum) % coreNum == 0U) 
    {   
        uint32_t blockLength = totalLengthAligned / coreNum;
        uint32_t tileNum, tileLength, lastTileLength;

        TilingParamsCalc(blockLength, alignNum, tileNum, tileLength, lastTileLength);

        tiling.set_blockLength(blockLength);
        tiling.set_tileNum(tileNum);
        tiling.set_tileLength(tileLength);
        tiling.set_lastTileLength(lastTileLength);
        tiling.set_isEvenCore(1U);
    } 
    else 
    {   
        uint32_t totalBlockCount = totalLengthAligned / alignNum;  
        uint32_t formerNum = totalBlockCount % coreNum;           
        uint32_t tailNum = coreNum - formerNum;               

        uint32_t formerBlockCount = (totalBlockCount + coreNum - 1) / coreNum;
        uint32_t tailBlockCount = totalBlockCount / coreNum;
        uint32_t formerLength = formerBlockCount * alignNum;
        uint32_t tailLength = tailBlockCount * alignNum;

        uint32_t fTileNum, fTileLen, fLastTileLen;
        TilingParamsCalc(formerLength, alignNum, fTileNum, fTileLen, fLastTileLen);
        
        uint32_t tTileNum, tTileLen, tLastTileLen;
        TilingParamsCalc(tailLength, alignNum, tTileNum, tTileLen, tLastTileLen);

        tiling.set_formerNum(formerNum);
        tiling.set_formerLength(formerLength);
        tiling.set_formerTileNum(fTileNum);
        tiling.set_formerTileLength(fTileLen);
        tiling.set_formerLastTileLength(fLastTileLen);

        tiling.set_tailNum(tailNum);
        tiling.set_tailLength(tailLength);
        tiling.set_tailTileNum(tTileNum);
        tiling.set_tailTileLength(tTileLen);
        tiling.set_tailLastTileLength(tLastTileLen);
        tiling.set_isEvenCore(0U);    
    }
   
    // 设置数据类型
    tiling.set_dataTypeStart(inputDataType_1);
    tiling.set_dataTypeEnd(inputDataType_2);
    tiling.set_dataTypeStep(inputDataType_3);

    tiling.set_totalLength(totalLength);
    tiling.set_totalLengthAligned(totalLengthAligned);          
    context->SetBlockDim(coreNum);
    printf("[TilingFunc] tilingparams success\n");
    printf("hellorange\n");
    auto rawTiling = context->GetRawTilingData();
    printf("[TilingFunc] GetRawTilingData success\n");
    if (rawTiling == nullptr) {
        printf("[TilingFunc] ERROR: RawTilingData is null\n");
        return ge::GRAPH_FAILED;
    }
    size_t capacity = rawTiling->GetCapacity();
    printf("[TilingFunc] RawTilingData capacity=%zu\n", capacity);
    uint32_t dataSize = tiling.GetDataSize();
    printf("[TilingFunc] RawTilingData capacity=%zu, tilingDataSize=%u\n", capacity, dataSize);
    if (dataSize > capacity) {
        printf("[TilingFunc] ERROR: tiling data overflow (dataSize=%u > capacity=%zu)\n", dataSize, capacity);
        return ge::GRAPH_FAILED;
    }

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), capacity);
    printf("[TilingFunc] SaveToBuffer success\n");
    rawTiling->SetDataSize(dataSize);
    printf("[TilingFunc] SetDataSize success\n");
    auto ws_num = context->GetWorkspaceNum();
    printf("[TilingFunc] workspace num=%zu\n", ws_num);
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    if (currentWorkspace == nullptr) {
        printf("[TilingFunc] ERROR: GetWorkspaceSizes returned null\n");
        return ge::GRAPH_FAILED;
    }
    currentWorkspace[0] = 0;
    printf("[TilingFunc] EXIT\n");
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    printf("[InferShape] ENTER\n");
    
    // 检查context指针是否为空
    if (context == nullptr) {
        printf("[InferShape] ERROR: InferShapeContext is null!\n");
        return ge::GRAPH_FAILED;
    }
    
    auto tensorStart = context->GetInputTensor(0);
    auto tensorEnd = context->GetInputTensor(1);
    auto tensorStep = context->GetInputTensor(2);
    
    if (tensorStart == nullptr || tensorEnd == nullptr || tensorStep == nullptr) {
        printf("[InferShape] ERROR: One or more input tensors are null!\n");
        return ge::GRAPH_FAILED;
    }
    
    // 计算Range的输出长度: ceil((end - start) / step)
    float startVal = 0.0f, endVal = 0.0f, stepVal = 1.0f;
    
    // 根据数据类型获取值
    const auto inputDataType_1 = context->GetInputDesc(0)->GetDataType();
    const auto inputDataType_2 = context->GetInputDesc(1)->GetDataType();
    const auto inputDataType_3 = context->GetInputDesc(2)->GetDataType();
    
    // 获取start值 - 参考LinSpaceD的实现，只支持float和int64
    switch(inputDataType_1) {
        case ge::DT_FLOAT:
            startVal = *reinterpret_cast<const float*>(tensorStart);
            break;
        case ge::DT_INT64:
            startVal = static_cast<float>(*reinterpret_cast<const int64_t*>(tensorStart));
            break;
        default:
            printf("[InferShape] ERROR: Unsupported start data type\n");
            return ge::GRAPH_FAILED;
    }
    
    // 获取end值
    switch(inputDataType_2) {
        case ge::DT_FLOAT:
            endVal = *reinterpret_cast<const float*>(tensorEnd);
            break;
        case ge::DT_INT64:
            endVal = static_cast<float>(*reinterpret_cast<const int64_t*>(tensorEnd));
            break;
        default:
            printf("[InferShape] ERROR: Unsupported end data type\n");
            return ge::GRAPH_FAILED;
    }
    
    // 获取step值
    switch(inputDataType_3) {
        case ge::DT_FLOAT:
            stepVal = *reinterpret_cast<const float*>(tensorStep);
            break;
        case ge::DT_INT64:
            stepVal = static_cast<float>(*reinterpret_cast<const int64_t*>(tensorStep));
            break;
        default:
            printf("[InferShape] ERROR: Unsupported step data type\n");
            return ge::GRAPH_FAILED;
    }
    
    // 计算输出长度: ceil(abs(end - start) / abs(step))
    float diff = std::abs(endVal - startVal);
    float stepAbs = std::abs(stepVal);
    int64_t output_length = static_cast<int64_t>(std::ceil(diff / stepAbs));

    gert::Shape* y_shape = context->GetOutputShape(0);
    if (y_shape == nullptr) {
        printf("[InferShape] ERROR: Output shape is null!\n");
        return ge::GRAPH_FAILED;
    }
    
    y_shape->SetDimNum(1);
    y_shape->SetDim(0, output_length);
    
    printf("[InferShape] EXIT\n");
    return ge::GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    printf("[InferDataType] ENTER\n");
    
    // 检查context指针是否为空
    if (context == nullptr) {
        printf("[InferDataType] ERROR: InferDataTypeContext is null!\n");
        return ge::GRAPH_FAILED;
    }
    
    const auto inputDataType_1 = context->GetInputDataType(0);
    const auto inputDataType_2 = context->GetInputDataType(1);
    const auto inputDataType_3 = context->GetInputDataType(2);
    
    bool isStartValid = (inputDataType_1 == ge::DT_INT64) || (inputDataType_1 == ge::DT_FLOAT);
    bool isEndValid = (inputDataType_2 == ge::DT_INT64) || (inputDataType_2 == ge::DT_FLOAT);
    bool isStepValid = (inputDataType_3 == ge::DT_INT64) || (inputDataType_3 == ge::DT_FLOAT);

    if (!isStartValid || !isEndValid || !isStepValid) {
        printf("[InferDataType] ERROR: Invalid input data types!\n");
        return ge::GRAPH_FAILED;
    }
    
    // 根据JSON配置，输出支持float和int64，但根据value_depend特性，我们输出float
    context->SetOutputDataType(0, ge::DT_FLOAT);
    
    printf("[InferDataType] EXIT\n");
    return ge::GRAPH_SUCCESS;
}
} // namespace ge


namespace ops {
class RangeCustom : public OpDef {
public:
    explicit RangeCustom(const char* name) : OpDef(name)
    {
        this->Input("start")
            .ParamType(REQUIRED)
            .ValueDepend(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("end")
        .ParamType(REQUIRED)
        .ValueDepend(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("step")
        .ParamType(REQUIRED)
        .ValueDepend(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(RangeCustom);
} // namespace ops
